import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import re

import torch
from transformers import BertForQuestionAnswering
from transformers import BertTokenizer

from transformers import pipeline
generator = pipeline('text-generation', model = 'gpt2')
generator.model.config.pad_token_id = generator.model.config.eos_token_id


def oracle(start_of_sentence):
    output = generator(start_of_sentence)
    answer = output[0]['generated_text']
    answer = re.sub('[^!?.]+$', '', answer)
    return answer

def main_cl():
    print("Ready")

    try:
        while True:
            question = input("> ")
            o = generator(question)
            answer = o[0]['generated_text']
            answer = re.sub('[^!?.]+$', '', answer)
            print(answer)
    except EOFError as e:
        pass


main_cl()
